#include "cmsisDAP.h"

#include <string.h>

struct pending_scan_result {
    /** Offset in bytes in the CMD_DAP_JTAG_SEQ response buffer. */
    unsigned first;
    /** Number of bits to read. */
    unsigned length;
    /** Location to store the result */
    uint8_t *buffer;
    /** Offset in the destination buffer */
    unsigned buffer_offset;
};

#define QUEUED_SEQ_BUF_LEN (cmsis_dap_handle->packet_size - 3)

/* CMD_CONNECT */
#define CONNECT_DEFAULT           0x00
#define CONNECT_SWD               0x01
#define CONNECT_JTAG              0x02

/* CMSIS-DAP JTAG sequence info masks */
/* Number of bits to clock through (0 means 64) */
#define DAP_JTAG_SEQ_TCK          0x3F
/* TMS will be set during the sequence if this bit is set */
#define DAP_JTAG_SEQ_TMS          0x40
/* TDO output will be captured if this bit is set */
#define DAP_JTAG_SEQ_TDO          0x80

#define USB_TIMEOUT       1000

/* CMSIS-DAP General Commands */
#define CMD_DAP_INFO              0x00
#define CMD_DAP_LED               0x01
#define CMD_DAP_CONNECT           0x02
#define CMD_DAP_DISCONNECT        0x03
#define CMD_DAP_WRITE_ABORT       0x08
#define CMD_DAP_DELAY             0x09
#define CMD_DAP_RESET_TARGET      0x0A

/* DAP Status Code */
#define DAP_OK                    0
#define DAP_ERROR                 0xFF

/* CMSIS-DAP Common SWD/JTAG Commands */
#define CMD_DAP_DELAY             0x09
#define CMD_DAP_SWJ_PINS          0x10
#define CMD_DAP_SWJ_CLOCK         0x11
#define CMD_DAP_SWJ_SEQ           0x12

/* CMSIS-DAP JTAG Commands */
#define CMD_DAP_JTAG_SEQ          0x14
#define CMD_DAP_JTAG_CONFIGURE    0x15
#define CMD_DAP_JTAG_IDCODE       0x16

cmsisDAP::cmsisDAP()
{
    memset(&_dap, 0, sizeof(_dap));
    queued_seq_buf = NULL;
    cmsis_dap_handle = &_dap;
}

bool cmsisDAP::init()
{
    queued_seq_count = 0;
    queued_seq_buf_end = 0;
    pending_scan_result_count = 0;
    queued_seq_tdo_ptr = 0;

    queued_seq_buf = (uint8_t*)malloc(1024);
    pending_scan_results = (struct pending_scan_result *)calloc(255, sizeof(struct pending_scan_result *));
    if (!queued_seq_buf || !pending_scan_results)
    {
        deinit();
        return false;
    }

    return true;
}

void cmsisDAP::deinit()
{
    if (queued_seq_buf)
        free(queued_seq_buf);
    queued_seq_buf = nullptr;
    if (pending_scan_results)
        free(pending_scan_results);
    pending_scan_results = nullptr;
}

int cmsisDAP::cmsis_dap_xfer(struct cmsis_dap *dap, int txlen)
{
    uint8_t current_cmd = dap->command[0];
    int retval = dap_write(dap, txlen, USB_TIMEOUT);
    if (retval < 0)
        return retval;

    /* get reply */
    retval = dap_read(dap, USB_TIMEOUT);
    if (retval < 0)
        return retval;

    uint8_t *resp = dap->response;
    if (resp[0] == DAP_ERROR) {
        LOG_ERROR("CMSIS-DAP command 0x%" PRIx8 " not implemented", current_cmd);
        return ERROR_NOT_IMPLEMENTED;
    }

    return 0;
}

int cmsisDAP::cmsis_dap_flush()
{
    if (!queued_seq_count)
        return 0;

    LOG_DEBUG_IO("Flushing %d queued sequences (%d bytes) with %d pending scan results to capture",
        queued_seq_count, queued_seq_buf_end, pending_scan_result_count);

    /* prepare CMSIS-DAP packet */
    uint8_t *command = cmsis_dap_handle->command;
    command[0] = CMD_DAP_JTAG_SEQ;
    command[1] = queued_seq_count;
    memcpy(&command[2], queued_seq_buf, queued_seq_buf_end);

#ifdef CMSIS_DAP_JTAG_DEBUG
    debug_parse_cmsis_buf(command, queued_seq_buf_end + 2);
#endif

    /* send command to USB device */
    int retval = cmsis_dap_xfer(cmsis_dap_handle, queued_seq_buf_end + 2);

    uint8_t *resp = cmsis_dap_handle->response;
    if (retval != ERROR_OK || resp[1] != DAP_OK) {
        LOG_ERROR("CMSIS-DAP command CMD_DAP_JTAG_SEQ failed.");
        return -1;
    }

#ifdef CMSIS_DAP_JTAG_DEBUG
    LOG_DEBUG_IO("USB response buf:");
    for (int c = 0; c < queued_seq_buf_end + 3; ++c)
        printf("%02X ", resp[c]);
    printf("\n");
#endif

    /* copy scan results into client buffers */
    for (int i = 0; i < pending_scan_result_count; ++i) {
        struct pending_scan_result *scan = &pending_scan_results[i];
        LOG_DEBUG_IO("Copying pending_scan_result %d/%d: %d bits from byte %d -> buffer + %d bits",
            i, pending_scan_result_count, scan->length, scan->first + 2, scan->buffer_offset);
#ifdef CMSIS_DAP_JTAG_DEBUG
        for (uint32_t b = 0; b < DIV_ROUND_UP(scan->length, 8); ++b)
            printf("%02X ", resp[2+scan->first+b]);
        printf("\n");
#endif
        bit_copy(scan->buffer, scan->buffer_offset, &resp[2 + scan->first], 0, scan->length);
    }

    /* reset */
    queued_seq_count = 0;
    queued_seq_buf_end = 0;
    queued_seq_tdo_ptr = 0;
    pending_scan_result_count = 0;

    return 0;
}

///////////////////////////////////////

int cmsisDAP::selectMode(enum DapMode mode)
{
    uint8_t m = CONNECT_DEFAULT;

    if (mode == DM_JTAG)
        m = CONNECT_JTAG;
    else if (mode == DM_SWD)
        m = CONNECT_SWD;

    return cmsis_dap_cmd_dap_connect(m);
}

int cmsisDAP::cmsis_dap_cmd_dap_connect(uint8_t mode)
{
    uint8_t *command = cmsis_dap_handle->command;

    command[0] = CMD_DAP_CONNECT;
    command[1] = mode;

    int retval = cmsis_dap_xfer(cmsis_dap_handle, 2);
    if (retval != ERROR_OK) {
        LOG_ERROR("CMSIS-DAP command CMD_CONNECT failed.");
        return ERROR_JTAG_DEVICE_ERROR;
    }

    if (cmsis_dap_handle->response[1] != mode) {
        LOG_ERROR("CMSIS-DAP failed to connect in mode (%d)", mode);
        return ERROR_JTAG_DEVICE_ERROR;
    }

    return ERROR_OK;
}

/////////////////////////////////////////////////////

int cmsisDAP::execute_queue(struct list_head *chead)
{
	int ret;
	struct jtag_command *cmd;

	__list_for_each_entry(cmd, chead, node, struct jtag_command)
	{
		ret = execute(cmd);
		if (ret != 0)
		{
			break;
		}
	}

    cmsis_dap_flush();

	return ret;
}

int cmsisDAP::execute(struct jtag_command *cmd)
{
	int ret = 0;

	switch (cmd->type)
	{
	case JTAG_SLEEP:
		cmsis_dap_flush();
		cmsis_dap_execute_sleep(cmd);
		break;
	case JTAG_TLR_RESET:
		cmsis_dap_flush();
		cmsis_dap_execute_tlr_reset(cmd);
		break;
	case JTAG_SCAN:
		cmsis_dap_execute_scan(cmd);
		break;
	case JTAG_PATHMOVE:
		cmsis_dap_execute_pathmove(cmd);
		break;
	case JTAG_RUNTEST:
		cmsis_dap_execute_runtest(cmd);
		break;
	case JTAG_STABLECLOCKS:
		cmsis_dap_execute_stableclocks(cmd);
		break;
	case JTAG_TMS:
		cmsis_dap_execute_tms(cmd);
		break;
	default:
		LOG_ERROR("BUG: unknown JTAG command type 0x%X encountered", cmd->type);
		break;
	}

	return ret;
}

void cmsisDAP::cmsis_dap_execute_tms(struct jtag_command *cmd)
{
	LOG_DEBUG_IO("TMS: %d bits", cmd->cmd.tms->num_bits);
	cmsis_dap_cmd_dap_swj_sequence(cmd->cmd.tms->num_bits, cmd->cmd.tms->bits);
}

void cmsisDAP::cmsis_dap_execute_stableclocks(struct jtag_command *cmd)
{
	LOG_DEBUG_IO("stableclocks %i cycles", cmd->cmd.runtest->num_cycles);
	cmsis_dap_stableclocks(cmd->cmd.runtest->num_cycles);
}

void cmsisDAP::cmsis_dap_execute_runtest(struct jtag_command *cmd)
{
	LOG_DEBUG_IO("runtest %i cycles, end in %i", cmd->cmd.runtest->num_cycles,
				 cmd->cmd.runtest->end_state);

	cmsis_dap_end_state(cmd->cmd.runtest->end_state);
	cmsis_dap_runtest(cmd->cmd.runtest->num_cycles);
}

void cmsisDAP::cmsis_dap_execute_sleep(struct jtag_command *cmd)
{
}

/* Set TMS high for five TCK clocks, to move the TAP to the Test-Logic-Reset state */
int cmsisDAP::cmsis_dap_execute_tlr_reset(struct jtag_command *cmd)
{
	LOG_INFO("cmsis-dap JTAG TLR_RESET");
	uint8_t seq = 0xff;

	int retval = cmsis_dap_cmd_dap_swj_sequence(8, &seq);
	if (retval == ERROR_OK)
		tap_set_state(TAP_RESET);
	return retval;
}

void cmsisDAP::cmsis_dap_execute_scan(struct jtag_command *cmd)
{
	LOG_DEBUG_IO("%s type:%d", cmd->cmd.scan->ir_scan ? "IRSCAN" : "DRSCAN",
				 jtag_scan_type(cmd->cmd.scan));

	/* Make sure there are no trailing fields with num_bits == 0, or the logic below will fail. */
	while (cmd->cmd.scan->num_fields > 0 && cmd->cmd.scan->fields[cmd->cmd.scan->num_fields - 1].num_bits == 0)
	{
		cmd->cmd.scan->num_fields--;
		LOG_DEBUG("discarding trailing empty field");
	}

	if (cmd->cmd.scan->num_fields == 0)
	{
		LOG_DEBUG("empty scan, doing nothing");
		return;
	}

	if (cmd->cmd.scan->ir_scan)
	{
		if (tap_get_state() != TAP_IRSHIFT)
		{
			cmsis_dap_end_state(TAP_IRSHIFT);
			cmsis_dap_state_move();
		}
	}
	else
	{
		if (tap_get_state() != TAP_DRSHIFT)
		{
			cmsis_dap_end_state(TAP_DRSHIFT);
			cmsis_dap_state_move();
		}
	}

	cmsis_dap_end_state(cmd->cmd.scan->end_state);

	struct scan_field *field = cmd->cmd.scan->fields;
	unsigned scan_size = 0;

	for (int i = 0; i < cmd->cmd.scan->num_fields; i++, field++)
	{
		scan_size += field->num_bits;
		LOG_DEBUG_IO("%s%s field %d/%d %d bits",
					 field->in_value ? "in" : "",
					 field->out_value ? "out" : "",
					 i,
					 cmd->cmd.scan->num_fields,
					 field->num_bits);

		if (i == cmd->cmd.scan->num_fields - 1 && tap_get_state() != tap_get_end_state())
		{
			LOG_DEBUG_IO("Last field and have to move out of SHIFT state");
			/* Last field, and we're leaving IRSHIFT/DRSHIFT. Clock last bit during tap
			 * movement. This last field can't have length zero, it was checked above. */
			cmsis_dap_add_jtag_sequence(
				field->num_bits - 1, /* number of bits to clock */
				field->out_value,	 /* output sequence */
				0,					 /* output offset */
				false,				 /* TMS low */
				field->in_value,
				0);

			/* Clock the last bit out, with TMS high */
			uint8_t last_bit = 0;
			if (field->out_value)
				bit_copy(&last_bit, 0, field->out_value, field->num_bits - 1, 1);
			cmsis_dap_add_jtag_sequence(
				1,
				&last_bit,
				0,
				true,
				field->in_value,
				field->num_bits - 1);
			tap_set_state(tap_state_transition(tap_get_state(), 1));

			/* Now clock one more cycle, with TMS low, to get us into a PAUSE state */
			cmsis_dap_add_jtag_sequence(
				1,
				&last_bit,
				0,
				false,
				NULL,
				0);
			tap_set_state(tap_state_transition(tap_get_state(), 0));
		}
		else
		{
			LOG_DEBUG_IO("Internal field, staying in SHIFT state afterwards");
			/* Clocking part of a sequence into DR or IR with TMS=0,
			   leaving TMS=0 at the end so we can continue later */
			cmsis_dap_add_jtag_sequence(
				field->num_bits,
				field->out_value,
				0,
				false,
				field->in_value,
				0);
		}
	}

	if (tap_get_state() != tap_get_end_state())
	{
		cmsis_dap_end_state(tap_get_end_state());
		cmsis_dap_state_move();
	}

	LOG_DEBUG_IO("%s scan, %i bits, end in %s",
				 (cmd->cmd.scan->ir_scan) ? "IR" : "DR", scan_size,
				 tap_state_name(tap_get_end_state()));
}

void cmsisDAP::cmsis_dap_execute_pathmove(struct jtag_command *cmd)
{
	cmsis_dap_pathmove(cmd->cmd.pathmove->num_states, cmd->cmd.pathmove->path);
}

//////////////////////////////////////

int cmsisDAP::_add_jtag_sequence(int s_len, const uint8_t *sequence, int s_offset,
					   bool tms, uint8_t *tdo_buffer, int tdo_buffer_offset)
{
	if (s_len == 0)
		return 0;

	int cmd_len = 1 + DIV_ROUND_UP(s_len, 8);
    if (queued_seq_count >= 255 || queued_seq_buf_end + cmd_len > QUEUED_SEQ_BUF_LEN)
		/* empty out the buffer */
		cmsis_dap_flush();

    ++queued_seq_count;

	/* control byte */
    queued_seq_buf[queued_seq_buf_end] =
		(tms ? DAP_JTAG_SEQ_TMS : 0) |
		(tdo_buffer ? DAP_JTAG_SEQ_TDO : 0) |
		(s_len == 64 ? 0 : s_len);

	if (sequence)
        bit_copy(&queued_seq_buf[queued_seq_buf_end + 1], 0, sequence, s_offset, s_len);
	else
        memset(&queued_seq_buf[queued_seq_buf_end + 1], 0, DIV_ROUND_UP(s_len, 8));

    queued_seq_buf_end += cmd_len;

	if (tdo_buffer)
	{
		struct pending_scan_result *scan = &pending_scan_results[pending_scan_result_count++];
		scan->first = queued_seq_tdo_ptr;
		queued_seq_tdo_ptr += DIV_ROUND_UP(s_len, 8);
		scan->length = s_len;
		scan->buffer = tdo_buffer;
		scan->buffer_offset = tdo_buffer_offset;
	}

	return 0;
}

int cmsisDAP::cmsis_dap_cmd_dap_swj_sequence(uint8_t s_len, const uint8_t *sequence)
{
    uint8_t *command = cmsis_dap_handle->command;

    command[0] = CMD_DAP_SWJ_SEQ;
    command[1] = s_len;
    bit_copy(&command[2], 0, sequence, 0, s_len);

    int retval = cmsis_dap_xfer(cmsis_dap_handle, 2 + DIV_ROUND_UP(s_len, 8));
    if (retval != ERROR_OK || cmsis_dap_handle->response[1] != DAP_OK)
        return ERROR_FAIL;

    return ERROR_OK;
}

void cmsisDAP::cmsis_dap_pathmove(int num_states, tap_state_t *path)
{
	uint8_t tms0 = 0x00;
	uint8_t tms1 = 0xff;

	for (int i = 0; i < num_states; i++)
	{
		if (path[i] == tap_state_transition(tap_get_state(), false))
			cmsis_dap_add_tms_sequence(&tms0, 1);
		else if (path[i] == tap_state_transition(tap_get_state(), true))
			cmsis_dap_add_tms_sequence(&tms1, 1);
		else
		{
			LOG_ERROR("BUG: %s -> %s isn't a valid TAP transition.",
					  tap_state_name(tap_get_state()), tap_state_name(path[i]));
			exit(-1);
		}

		tap_set_state(path[i]);
	}

	cmsis_dap_end_state(tap_get_state());
}

/* queue a sequence of bits to clock out TMS, executing if the buffer is full */
void cmsisDAP::cmsis_dap_add_tms_sequence(const uint8_t *sequence, int s_len)
{
	LOG_DEBUG_IO("%d bits: %02X", s_len, *sequence);
	/* we use a series of CMD_DAP_JTAG_SEQ commands to toggle TMS,
	   because even though it seems ridiculously inefficient, it
	   allows us to combine TMS and scan sequences into the same
	   USB packet. */
	/* TODO: combine runs of the same tms value */
	for (int i = 0; i < s_len; ++i)
	{
		bool bit = (sequence[i / 8] & (1 << (i % 8))) != 0;
		cmsis_dap_add_jtag_sequence(1, NULL, 0, bit, NULL, 0);
	}
}

/* queue a sequence of bits to clock out TDI / in TDO, executing if the buffer is full.
 *
 * sequence=NULL means clock out zeros on TDI
 * tdo_buffer=NULL means don't capture TDO
 */
void cmsisDAP::cmsis_dap_add_jtag_sequence(int s_len, const uint8_t *sequence, int s_offset,
										   bool tms, uint8_t *tdo_buffer, int tdo_buffer_offset)
{
	for (int offset = 0; offset < s_len; offset += 64)
	{
		int len = s_len - offset;

		if (len > 64)
			len = 64;

		_add_jtag_sequence(
			len,
			sequence,
			s_offset + offset,
			tms,
			tdo_buffer,
			!tdo_buffer ? 0 : (tdo_buffer_offset + offset));
	}
}

/* Set new end state */
void cmsisDAP::cmsis_dap_end_state(tap_state_t state)
{
	if (tap_is_state_stable(state))
		tap_set_end_state(state);
	else
	{
		LOG_ERROR("BUG: %i is not a valid end state", state);
		exit(-1);
	}
}

/* Move to the end state by queuing a sequence to clock into TMS */
void cmsisDAP::cmsis_dap_state_move(void)
{
	uint8_t tms_scan = tap_get_tms_path(tap_get_state(), tap_get_end_state());
	uint8_t tms_scan_bits = tap_get_tms_path_len(tap_get_state(), tap_get_end_state());

	LOG_DEBUG_IO("state move from %s to %s: %d clocks, %02X on tms",
				 tap_state_name(tap_get_state()), tap_state_name(tap_get_end_state()),
				 tms_scan_bits, tms_scan);
	cmsis_dap_add_tms_sequence(&tms_scan, tms_scan_bits);

	tap_set_state(tap_get_end_state());
}

void cmsisDAP::cmsis_dap_stableclocks(int num_cycles)
{
	uint8_t tms = tap_get_state() == TAP_RESET;
	/* TODO: Perform optimizations? */
	/* Execute num_cycles. */
	for (int i = 0; i < num_cycles; i++)
		cmsis_dap_add_tms_sequence(&tms, 1);
}

void cmsisDAP::cmsis_dap_runtest(int num_cycles)
{
	tap_state_t saved_end_state = tap_get_end_state();

	/* Only do a state_move when we're not already in IDLE. */
	if (tap_get_state() != TAP_IDLE)
	{
		cmsis_dap_end_state(TAP_IDLE);
		cmsis_dap_state_move();
	}
	cmsis_dap_stableclocks(num_cycles);

	/* Finish in end_state. */
	cmsis_dap_end_state(saved_end_state);

	if (tap_get_state() != tap_get_end_state())
		cmsis_dap_state_move();
}
